
"""
EPS2到SAC格式转换模块

该模块实现了将EPS2格式数据转换为SAC格式的功能。
EPS2是一种地震数据格式，通常包含地震仪记录的三个分量数据。

主要功能：
1. 解析EPS2数据文件
2. 提取元数据信息
3. 将数据转换为SAC格式
4. 保存SAC文件

使用方法：
    from eps2_to_sac import convert_eps2_to_sac
    
    # 转换单个EPS2文件
    sac_files = convert_eps2_to_sac("path/to/eps2_file.012", "output_dir")
    
    # 批量转换目录中的EPS2文件
    sac_files = convert_eps2_directory_to_sac("input_dir", "output_dir")
"""

import os
import struct
import numpy as np
from pathlib import Path
from datetime import datetime
from typing import List, Tuple, Dict, Optional, Union

from obspy import Trace, Stream
from obspy.core import AttribDict
from obspy.core.utcdatetime import UTCDateTime


def parse_eps2_header(file_path: Union[str, Path]) -> Dict:
    """
    解析EPS2文件头，提取元数据信息
    
    参数:
        file_path: EPS2文件路径
        
    返回:
        包含元数据的字典
    """
    file_path = Path(file_path)
    metadata = {}

    # 首先，尝试从EPS2文件本身的72字节头部读取信息
    try:
        with open(file_path, 'rb') as f:
            header = f.read(72)
            if len(header) < 72:
                # 文件头不完整，可能无法解析
                pass

            # 根据文档解析文件头
            # 0-8: 仪器编号 (假定为字符串)
            metadata['instrument_id'] = header[0:9].decode('ansic', errors='ignore').strip('\x00')
            
            # 9-10: 采样率 (unsigned short, little-endian)
            metadata['sampling_rate'] = struct.unpack('<H', header[9:11])[0]
            
            # 11-18: 纬度
            lat_deg = struct.unpack('<I', header[11:15])[0]
            lat_frac = struct.unpack('<I', header[15:19])[0]
            metadata['latitude'] = lat_deg + lat_frac / 1e6 # 假设是百万分之一度

            # 19-26: 经度
            lon_deg = struct.unpack('<I', header[19:23])[0]
            lon_frac = struct.unpack('<I', header[23:27])[0]
            metadata['longitude'] = lon_deg + lon_frac / 1e6 # 假设是百万分之一度

            # 27-28: 仪器增益
            metadata['gain'] = struct.unpack('<H', header[27:29])[0]
            
            # 36-47: 台站ID
            metadata['station_id'] = header[36:48].decode('ansic', errors='ignore').strip('\x00')
            
            # 48-59: 站点名
            metadata['station_name'] = header[48:60].decode('ansic', errors='ignore').strip('\x00')

    except Exception:
        # 如果从文件头读取失败，继续尝试其他方法
        pass

    # 从文件名中提取信息
    stem = file_path.stem
    if len(stem) == 8:
        try:
            hex_value = int(stem, 16)
            metadata['file_id'] = hex_value
            metadata['file_id_hex'] = stem
        except ValueError:
            metadata['file_id'] = None
            metadata['file_id_hex'] = stem
    
    # 尝试从DATAFILE.LST中获取更多信息
    lst_file = file_path.parent / "DATAFILE.LST"
    if lst_file.exists():
        try:
            with open(lst_file, 'r', encoding='utf-8', errors='ignore') as f:
                for line in f:
                    parts = line.strip().split()
                    # 文件名可能包含 "1:" 前缀
                    if len(parts) > 3 and stem in parts[3]:
                        try:
                            # 格式: 0001 2019-01-10 01:38:54.0122 1:4C5419B6.012 100sps
                            datetime_str = f"{parts[1]} {parts[2]}"
                            # LST文件中的时间可能包含毫秒或更精确
                            if '.' in parts[2]:
                                metadata['datetime'] = datetime.strptime(datetime_str, "%Y-%m-%d %H:%M:%S.%f")
                            else:
                                metadata['datetime'] = datetime.strptime(datetime_str, "%Y-%m-%d %H:%M:%S")

                            # 更新采样率（如果LST文件中有）
                            for part in parts:
                                if 'sps' in part.lower():
                                    metadata['sampling_rate'] = int(part.lower().replace('sps', ''))
                                    break
                            break
                        except (ValueError, IndexError):
                            continue
        except Exception:
            pass
    
    # 尝试从LOG文件中获取更多信息
    log_files = list(file_path.parent.glob("*.LOG"))
    if log_files:
        log_file = log_files[0]
        try:
            with open(log_file, 'r', encoding='utf-8', errors='ignore') as f:
                for line in f:
                    if 'Pos=' in line:
                        try:
                            pos_str = line.split('Pos=')[1].split('.')[0]
                            # 简单的经纬度解析
                            if 'N' in pos_str and ('E' in pos_str or 'W' in pos_str):
                                metadata['position_log'] = pos_str
                        except IndexError:
                            pass
                    if not metadata.get('datetime') and line.startswith('@'):
                         # 格式: @2014-07-17 02:56:50 ...
                        try:
                            dt_str = line.split('@')[1].strip().split(' :')[0]
                            metadata['datetime'] = datetime.strptime(dt_str, "%Y-%m-%d %H:%M:%S")
                        except (IndexError, ValueError):
                            pass

        except Exception:
            pass
            
    return metadata


def read_eps2_data(file_path: Union[str, Path]) -> Tuple[np.ndarray, Dict]:
    """
    读取EPS2数据文件
    
    参数:
        file_path: EPS2文件路径
        
    返回:
        (数据数组, 元数据字典)
    """
    file_path = Path(file_path)
    
    # 解析文件头获取元数据
    metadata = parse_eps2_header(file_path)
    
    with open(file_path, 'rb') as f:
        # 跳过72字节的文件头
        f.seek(72)
        
        # 读取剩余的二进制数据
        data_bytes = f.read()
        
        # 数据存储为12B一组，每个分量4B的int型，小端序
        # 实际数据是原始COUNT值左移四位
        # 因此，我们需要右移四位来获取真实值
        
        # 确保数据长度是12的倍数 (Z, X, Y)
        num_bytes_per_group = 12
        if len(data_bytes) % num_bytes_per_group != 0:
            extra_bytes = len(data_bytes) % num_bytes_per_group
            data_bytes = data_bytes[:-extra_bytes]

        # 将字节流解析为32位小端整数
        data = np.frombuffer(data_bytes, dtype='<i4')
        
        # 右移4位以获取真实的COUNT值
        data = data >> 4
        
        # 数据是Z, X, Y三分量交织的
        num_points = len(data) // 3
        
        # 将数据重塑为 (num_points, 3)
        data = data.reshape(num_points, 3)
        
        # 根据文档，顺序是 Z, X, Y
        # SAC标准通常是 Z, N, E
        # 这里我们暂时保持 Z, X, Y，可以在创建Trace时再调整
        metadata['npts'] = num_points
        metadata['components'] = ['Z', 'X', 'Y']

    return data, metadata

def create_sac_trace(data: np.ndarray, metadata: Dict, component: str = 'Z') -> Trace:
    """
    创建SAC道对象
    
    参数:
        data: 数据数组
        metadata: 元数据字典
        component: 分量标识 ('X', 'Y', 'Z')
        
    返回:
        SAC道对象
    """
    # 创建SAC头
    sac_header = AttribDict()
    
    # 设置基本参数
    sac_header.npts = len(data)
    sac_header.delta = 1.0 / metadata.get('sampling_rate', 100)
    
    # 设置开始时间
    if 'datetime' in metadata:
        start_time = UTCDateTime(metadata['datetime'])
    elif 'log_datetime' in metadata:
        start_time = UTCDateTime(metadata['log_datetime'])
    else:
        # 如果没有时间信息，使用当前时间
        start_time = UTCDateTime()
    
    sac_header.starttime = start_time
    
    # 设置分量
    if component == 'X':
        sac_header.kcmpnm = 'E' # 东西向
        sac_header.cmpaz = 90.0
        sac_header.cmpinc = 90.0
    elif component == 'Y':
        sac_header.kcmpnm = 'N' # 南北向
        sac_header.cmpaz = 0.0
        sac_header.cmpinc = 90.0
    else:  # Z
        sac_header.kcmpnm = 'Z' # 垂直
        sac_header.cmpaz = 0.0
        sac_header.cmpinc = 0.0
    
    # 设置台站信息
    if 'station_name' in metadata and metadata['station_name']:
        sac_header.kstnm = metadata['station_name']
    elif 'station_id' in metadata and metadata['station_id']:
        sac_header.kstnm = metadata['station_id']
    elif 'file_id_hex' in metadata:
        sac_header.kstnm = metadata['file_id_hex'][:4]

    if 'latitude' in metadata:
        sac_header.stla = metadata['latitude']
    if 'longitude' in metadata:
        sac_header.stlo = metadata['longitude']

    # 创建SAC道
    trace = Trace(data=data, header=sac_header)
    
    return trace


def convert_eps2_to_sac(input_path: Union[str, Path], 
                        output_dir: Union[str, Path],
                        prefix: Optional[str] = None) -> List[Path]:
    """
    将EPS2文件转换为SAC格式
    
    参数:
        input_path: 输入EPS2文件路径
        output_dir: 输出目录
        prefix: 输出文件名前缀（可选）
        
    返回:
        生成的SAC文件路径列表
    """
    input_path = Path(input_path)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        data, metadata = read_eps2_data(input_path)
    except Exception as e:
        raise ValueError(f"读取EPS2文件失败: {e}")
    
    components = metadata.get('components', ['Z', 'X', 'Y'])
    
    # data.shape is (npts, 3)
    z_data = data[:, 0]
    x_data = data[:, 1]
    y_data = data[:, 2]
    data_arrays = [z_data, x_data, y_data]

    if prefix is None:
        prefix = input_path.stem
    
    subdir_info = ""
    if input_path.parent.name != output_dir.name and input_path.parent.name != ".":
        subdir_info = f"{input_path.parent.name}_"

    sac_files = []
    for comp, data_array in zip(components, data_arrays):
        trace = create_sac_trace(data_array, metadata, comp)
        
        output_name = f"{subdir_info}{prefix}_{comp}.sac"
        output_path = output_dir / output_name
        
        trace.write(str(output_path), format='SAC')
        sac_files.append(output_path)
    
    return sac_files

def convert_eps2_directory_to_sac(input_dir: Union[str, Path],
                                 output_dir: Union[str, Path],
                                 recursive: bool = True) -> List[Path]:
    """
    批量转换目录中的EPS2文件为SAC格式
    
    参数:
        input_dir: 输入目录
        output_dir: 输出目录
        recursive: 是否递归处理子目录
        
    返回:
        生成的SAC文件路径列表
    """
    input_dir = Path(input_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    if recursive:
        eps2_files = list(input_dir.rglob("*.012"))
    else:
        eps2_files = list(input_dir.glob("*.012"))
    
    if not eps2_files:
        raise ValueError(f"在目录 {input_dir} 中未找到EPS2文件 (*.012)")
    
    all_sac_files = []
    for eps2_file in eps2_files:
        try:
            relative_path = eps2_file.relative_to(input_dir)
            file_output_dir = output_dir / relative_path.parent
            file_output_dir.mkdir(parents=True, exist_ok=True)
            
            sac_files = convert_eps2_to_sac(eps2_file, file_output_dir)
            all_sac_files.extend(sac_files)
        except Exception as e:
            print(f"转换文件 {eps2_file} 失败: {e}")
    
    return all_sac_files


if __name__ == '__main__':
    import sys
    
    if len(sys.argv) < 3:
        print("用法: python eps2_to_sac.py <输入EPS2文件或目录> <输出目录>")
        sys.exit(1)
    
    input_path = sys.argv[1]
    output_dir = sys.argv[2]
    
    if os.path.isfile(input_path):
        sac_files = convert_eps2_to_sac(input_path, output_dir)
        print(f"已转换EPS2文件为SAC格式，生成 {len(sac_files)} 个文件:")
        for f in sac_files:
            print(f"  {f}")
    elif os.path.isdir(input_path):
        sac_files = convert_eps2_directory_to_sac(input_path, output_dir)
        print(f"已批量转换EPS2文件为SAC格式，生成 {len(sac_files)} 个文件:")
        for f in sac_files:
            print(f"  {f}")
    else:
        print(f"错误: 输入路径 {input_path} 不是文件或目录")
        sys.exit(1)
